#!/usr/bin/python

# Using some super futuristic techniques involving the Matrix, I have devised a way to find TMS repeats in the most fashionable way

import blast
import hmmtop
import hmmgap
import tcdb
import tempfile
import StringIO
import os
import re
from Bio import SeqIO
from Bio import AlignIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC
from Bio.Blast.Applications import NcbiblastpCommandline
from Bio.Emboss.Applications import NeedleCommandline
from Bio.Blast import NCBIXML
from pysqlite2 import dbapi2 as sqlite
import urllib

class build:
	
	directory = "%s/desktop/Data/DB.2.A.1/" %(os.getenv('HOME'))
	dbfile = "mfs.db"

	def __init__(self):
		self.family = False
		self.accessions = []
		self.fasta_file = False
		self.tms = []
		self.fasta_dict = {}
		self.blast_status = False
		self.split = 2
		return
		
	def __call__(self):
		self.generate_fastas()
		self.hmmtop()
		self.build_tms_db()
		self.split_database(self.split)
		self.build_blast_db()
		return
		
	def generate_fastas(self):
		pb = blast.tools()
		for acc in self.accessions:
			print acc
			pb.psiblast(acc)
		pb.build_xml()
		pb.maketable(0.7)
		self.blast_status = pb.status
		self.fasta_file = pb.fasta_file
		return
		
	def hmmtop(self):
		hm = hmmtop.tools()
		hm.add_library("SF", self.fasta_file.name)
		hm.scan_libraries()
		self.tms = hm.results['SF']
		return
		
	def get_groups(self,keys):
		db = {}
		for size in keys:
			for row in keys:
				end = row+(size-1)
				if end > len(keys):
					break
				db.setdefault(size,[]).append(range(row,end+1))
		return db
		
	def define_family(self):
		url = "http://www.tcdb.org/projectv/sft.php?fam=%s" %(self.family)
 		response = urllib.urlopen(url)
		acc = response.read().strip()
		self.accessions = [i.strip() for i in acc.split(' ')]
		
	def build_tms_db(self):
		fastas = SeqIO.parse(self.fasta_file,"fasta")
		self.fasta_dict = SeqIO.to_dict(fastas)
		for fasta in self.fasta_dict.values():
			if fasta.name not in self.tms.keys():
				continue
			groups = self.get_groups(self.tms[fasta.name].keys())
			for group, tmss in groups.items():
				database = open(self.directory+str(group)+"_"+self.family+"_TMDB.faa","ab+")
				for tms in tmss:
					string = ''.join([str(self.fetch_tms(fasta.name,i)) for i in tms])
					count = ",".join([str(i) for i in tms])
					summary = "|| TMS#(%s) of %i || TC:%s |" %(count,len(self.tms[fasta.name].keys()), self.family)
					record = SeqRecord(Seq(string,IUPAC.protein),id=fasta.id, name=fasta.name, description=summary+fasta.description)
					SeqIO.write(record,database,'fasta')
				database.close()
		return
		
	def split_database(self,tms):
		dbfile = "%s%i_%s_TMDB.faa" %(self.directory,tms,self.family)
		bigdb = open(dbfile)
		fastas = SeqIO.parse(bigdb,'fasta')
		handle = open(self.directory+'buffer','w')
		for fasta in fastas:
			info = self.grab_header_info(fasta.description)
			newfile = "%ssplit/%i.%s.%s.faa" %(self.directory,info[1],'-'.join([str(i) for i in info[0]]),self.family)
			if handle.name is not newfile:
				handle.close()
				handle = open(newfile,"ab+")
			SeqIO.write(fasta,handle,'fasta')
		
		
	def build_blast_db(self):
		files = os.listdir(self.directory+'split/')
		fastas = [i for i in files if re.search('faa$',i)]
		for fasta in fastas:
			cmd = 'makeblastdb -in %s -out %s -dbtype prot' %(self.directory+'split/'+fasta,self.directory+'split/'+fasta)
			os.system(cmd)
		return
			
	def fetch_tms(self,name,tms):
		(start,end) = self.tms[name][tms]
		tms = self.fasta_dict[name].seq[start-1:end]
		return tms
		
	def grab_header_info(self,header):
		total = re.compile('of (\d+) \|\|')
		these = re.compile('TMS#\(([0-9,]+)\)')
		tmtotal = total.search(header).groups()[0]
		tmrange =these.search(header).groups()[0]
		if tmrange.isdigit() is False:
			tmrange=[int(i) for i in tmrange.split(",")]
		else:
			tmrange=[int(tmrange)]
		return [tmrange, int(tmtotal)]
		
	def get_symbol(self,header):
		pattern = re.compile('(\w+)\s+\|\|')
		symbol = pattern.search(header).groups()[0]
		sym = re.sub('\d+','',symbol)
		return symbol, sym
		
	def make_sql(self):
		self.connect()
		cursor = self.connection.cursor()
		query ='\
		CREATE TABLE "results" (\
			 "tms_total" integer(10,0) NOT NULL,\
			 "subject_tms" text(32,0) NOT NULL,\
			 "target_tms" text(32,0) NOT NULL,\
			 "subject_header" text(200,0) NOT NULL,\
			 "target_header" text(200,0) NOT NULL,\
			 "smatch" text(200,0) NOT NULL,\
			 "mmatch" text(200,0) NOT NULL,\
			 "tmatch" text(200,0) NOT NULL,\
			 "e" integer(32,50) NOT NULL,\
			 "bit" integer(32,0) NOT NULL,\
			 "len" integer(32,0) NOT NULL,\
			 "family" text(32,0) NOT NULL,\
			 "repeats" integer(10,0) NOT NULL,\
			 "subject_symbol" text(32,0) NOT NULL,\
			 "target_symbol" text(32,0) NOT NULL,\
			 "subject_sym" text(32,0) NOT NULL,\
			 "target_sym" text(32,0) NOT NULL,\
			 "id" integer NOT NULL,\
			 PRIMARY KEY("id")\
		);'
		cursor.execute(query)
		self.connection.commit()
		cursor.close()
		self.connection.close()
		
	def connect(self):
		self.connection = sqlite.connect(self.directory+self.dbfile)
		return
		

class search(build):
	
	def __init__(self, family, tms):
		self.outfmt = 5
		self.threads = 16
		self.expect = None # Minimum e-threshold
		self.out = tempfile.NamedTemporaryFile()
		self.records = False
		self.family = family
		self.tms = tms
		self.group = False # Only BLAST specific TMS Group. Defiend by Total TMS number
		self.connect()
		return
		
	def __call__(self):
		self.blast_targets()
		self.connection.close()
		return
		
	def parse_results(self):
		self.out.seek(0)
		self.records = NCBIXML.parse(self.out)
		return
						
	def blast_targets(self):
		targets= {}
		pblast = NcbiblastpCommandline()
		pblast.outfmt=5
		pblast.out=self.out.name
		pblast.num_threads=self.threads
		fastas = [i for i in os.listdir(self.directory+'split/') if re.search(self.family+'\.faa$',i)]
		[targets.setdefault(int(k),[]).append(v) for k,v in [self.parse_filename(i) for i in fastas]]
		
		for total,tmss in targets.items():
			if self.group is not False and self.group is not total:
				continue
			tmss.sort()
			for i,tms in enumerate(tmss):
				targets = tmss[i+self.tms:]
				if len(targets) is 0:
					continue
				targets.sort()
				subject_file = "%ssplit/%i.%s.%s.faa" %(self.directory,total,'-'.join([str(i) for i in tms]),self.family)
				for target in targets:
					target_file = "%ssplit/%i.%s.%s.faa" %(self.directory,total,'-'.join([str(i) for i in target]),self.family)
					pblast.query = subject_file
					pblast.db = target_file
					print "BLASTING TM-Group %i with (%s) vs (%s) in family: %s" %(total,','.join([str(i) for i in tms]),','.join([str(i) for i in target]), self.family)
					pblast()
					self.parse_results()
					self.collect_results(total)
					
	def collect_results(self,total):
		db = self.connection.cursor()
		for record in self.records:
			subject_tms=','.join([str(i) for i in self.grab_header_info(record.query)[0]])
			for alignment in record.alignments:
				target_tms=','.join([str(i) for i in self.grab_header_info(alignment.title)[0]])
				for hsp in alignment.hsps:
					subject_symbol = self.get_symbol(record.query)
					target_symbol = self.get_symbol(alignment.title)
					insert = 'insert into "main"."results" ( "tms_total", "subject_tms", "target_tms", "subject_header", "target_header",\
					"smatch", "mmatch", "tmatch", "e", "bit", "len", "family", "repeats","subject_symbol","target_symbol", "subject_sym", "target_sym") \
					values ( ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)'
					db.execute(insert,(total, subject_tms, target_tms, record.query, alignment.title, hsp.query, hsp.match, hsp.sbjct, hsp.expect, hsp.bits, alignment.length, self.family, self.tms, subject_symbol[0], target_symbol[0], subject_symbol[1], target_symbol[1]))
					self.connection.commit()
		db.close()
		return
							
	def parse_filename(self,name):
		parts = re.compile('(\d+)\.([-\d]+)\.([0-9A-Za-z.]+)\.faa')
		res = parts.search(name).groups()
		tms = [int(i) for i in res[1].split('-')]
		return res[0], tms
		
		
		
class extract(search,build):
	
	def __init__(self):
		# Set extraction methods
		self.e = "10000"
		self.identity = False
		self.simmilar = False
		self.max = 400
		self.length = 55
		self.binary = True
		self.repeats = 3
		self.groups = False
		self.family = "SPU10"
		self.dbfile="spu10.db"
		self.tcdb = False
		self.custom = False
		self.connect()
	
	def __call__(self):
		self.search_db()
		self.build_fasta()
		self.generate_report()
		
	def search_db(self):
		(ident,simmilar,length,groups,family,custom) = ['','','','','', '']
		if self.identity is True:
			ident = "AND subject_symbol=target_symbol"
		if self.simmilar is True:
			simmilar="AND subject_sym=target_sym"
		if self.length is not False:
			length="AND length(mmatch)>=%i"%(self.length)
		if self.groups is not False:
			groups="AND tms_total=%i"%(self.groups)
		if self.family is not False:
			family="AND `family`='%s'"%(self.family)
		if self.custom is not False:
			custom = 'AND %s' %(self.custom)
		query = "SELECT * FROM results WHERE e < '%s' AND repeats ='%i' %s %s %s %s %s %s ORDER BY e ASC LIMIT %i" %(self.e,self.repeats,ident,simmilar,length,groups, family, custom, self.max)
		print query
		db = self.connection.cursor()
		db.execute(query)
		self.results=db.fetchall()
	
	def build_fasta(self):
		subjects = list(set([self.extract_gi(i[3]) for i in self.results]))
		targets = list(set( [self.extract_gi(i[4]) for i in self.results]))
		gis = subjects + list(set(targets)-set(subjects))
		if self.tcdb is True:
			self.fasta_file = tcdb.api().acc2fasta(gis)
			return
		psi=blast.tools()
		psi.gis=gis
		self.fasta_file = psi.build_raw_fasta()
		
	def extract_gi(self,header):
		gix=re.compile('(gi|TC-DB)\|(\w+)')
		gi = gix.search(header).groups()
		return gi[1]
		
	def generate_report(self,filename="report.html"):
		handle = open(self.directory+filename,"wb+")
		fastas = SeqIO.parse(self.fasta_file,"fasta")
		fastas = SeqIO.to_dict(fastas)
		hmt = hmmtop.tools()
		hmt.add_library("SFT",self.fasta_file.name)
		hmt.scan_libraries()
		hmg = hmmgap.annotate()
		hmg.hmmtop=hmt.results['SFT']
		# Preform pairwise
		needle = NeedleCommandline()
		needle.gapopen = 8
		needle.gapextend =2
		needle.outfile ='stdout'
		# Open template file and get replacement module
		template = open("template.html",'r')
		module = re.search('\{row\}(.+)\{/row\}',template.read(),re.DOTALL).groups()[0]
		html=[]
		#handle.write('<pre>')
		for row in self.results:
			subject=fastas[str(self.extract_gi(row[3]))].seq
			target=fastas[str(self.extract_gi(row[4]))].seq
			needle.asequence="asis:"+str(subject)
			needle.bsequence="asis:"+str(target)
			res = needle()
			out=tempfile.NamedTemporaryFile()
			out.write(res[0])
			out.seek(0)
			alignment = AlignIO.read(out,'emboss')
			out.close()
			subject_file = tempfile.NamedTemporaryFile()
			SeqIO.write(fastas[str(self.extract_gi(row[3]))],subject_file,'fasta')
			subject_file.seek(0)
			target_file = tempfile.NamedTemporaryFile()
			SeqIO.write(fastas[str(self.extract_gi(row[4]))],target_file,'fasta')
			target_file.seek(0)
			hmg(subject_file.name,str(alignment[0].seq))
			subject_annotation=hmg.annotation
			hmg(target_file.name,str(alignment[1].seq))
			target_annotation=hmg.annotation
			entry = module.replace("{e}",str(row[8]))
			entry = entry.replace("{match_len}",str(len(row[6])))
			entry = entry.replace("{subject_tms}",row[1])
			entry = entry.replace("{target_tms}",row[2])
			entry = entry.replace("{subject_header}",row[3])
			entry = entry.replace("{target_header}",row[4])
			entry = entry.replace("{aln_s}",row[5])
			entry = entry.replace("{aln_m}",row[6])
			entry = entry.replace("{aln_t}",row[7])
			entry = entry.replace("{subject_annotation}",subject_annotation)
			entry = entry.replace("{subject_sequence}",str(alignment[0].seq))
			entry = entry.replace("{target_annotation}",target_annotation)
			entry = entry.replace("{target_sequence}",str(alignment[1].seq))
			entry = entry.replace("{family}",row[11])
			entry = entry.replace("{subject_symbol}",row[13])
			entry = entry.replace("{target_symbol}",row[14])
			entry = entry.replace("{tms_total}",str(row[0]))
			entry = entry.replace("{subject_fasta}",str(fastas[str(self.extract_gi(row[3]))].seq))
			entry = entry.replace("{target_fasta}",str(fastas[str(self.extract_gi(row[4]))].seq))
			html.append(entry)
		template.seek(0)
		mytemplate = template.read()
		mytemplate = mytemplate.replace("{tms}",str(self.repeats))
		mytemplate = mytemplate.replace("{show_total}",str(self.max))
		mytemplate = mytemplate.replace("{family}",self.family)
		output = re.sub('\{row\}.+\{/row\}',"\n".join(html),mytemplate,flags=re.DOTALL)
		handle.write(output)